# ----------------------------------------------------------------------------------
# Microsoft Developer & Platform Evangelism
# 
# Copyright (c) Microsoft Corporation. All rights reserved.
# 
# THIS CODE AND INFORMATION ARE PROVIDED "AS IS" WITHOUT WARRANTY OF ANY KIND, 
# EITHER EXPRESSED OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE IMPLIED WARRANTIES 
# OF MERCHANTABILITY AND/OR FITNESS FOR A PARTICULAR PURPOSE.
# ----------------------------------------------------------------------------------
# The example companies, organizations, products, domain names,
# e-mail addresses, logos, people, places, and events depicted
# herein are fictitious.  No association with any real company,
# organization, product, domain name, email address, logo, person,
# places, or events is intended or should be inferred.
# ----------------------------------------------------------------------------------
#
# THIS IS SAMPLE CODE. BUGS ARE EXPECTED. SECURITY IS NOT.
#
# This is the listener agent that has to run on the worker nodes.
# NOTE: it runs whatever executable is submitted. There is NO SAFETY CHECK.
# We strongly suggest that you implement such checks before use in production.
#

from datetime import *
from subprocess import *
from ast import *
from azure.storage import *
from azure.servicebus import *
from socket import *
from time import *
from azure import *
from os import *
from base64 import *
from shlex import *
from random import *


#name of node executing job
node=gethostname()
#name of queue in
qname = 'inqueue'
#name of job status table
jobtable = 'jobtable'
#storage account
storage_account='your storage account'
storage_key='your storage key'
#name of command queue (actually a service bus topic)
cqname = 'cqueue'
#name of result queue (Running, Completed, Failed, actually a service bus topic)
rqname = 'outqueue'
#service bus namespace
sbname='gerespy'
#service bus key
sbkey='your service bus key'
#service bus issuer
sbissuer='owner'

#open connections to azure storage queue, table, command and results topics
queue_service = QueueService(account_name=storage_account, account_key=storage_key)
table_service = TableService(account_name=storage_account, account_key=storage_key)
bus_service = ServiceBusService(service_namespace=sbname, account_key=sbkey, issuer=sbissuer)
     
#max idle time before shutdown
maxidle = timedelta(minutes=10)

#maximum job execution time in seconds
#the default visibility timeout for a message in a queue is 30 seconds. It may take longer than that for job to run.
#also, the time to live for a message in a queue is 7 days by default. We did not change that, but
#if a job is not executed within 7 days, it will drop off the queue.
maxtime = 5*60


#update job status and send a message
def update_status (table, node, id, task, newstatus):      
        print 'Current status: ', task.Status
        task.Status=newstatus
        task.TimeUpdated=datetime.now().isoformat()
        msg=Message(newstatus,custom_properties={'JobID': id})
        #update status table
        table_service.insert_or_replace_entity(table_name=table, partition_key=node, row_key=id, entity=task)
        #send message with status update
        bus_service.send_topic_message(rqname,msg)
        print 'New status: ', task.Status

#read job table entry for specific job id
def jobquery(jobtable,jobid):
    filter="RowKey eq '"+jobid+"'"
    print filter
    tasks=table_service.query_entities(jobtable,filter)
    print '%-10s\t%-16s\t%-10s\t%-26s\t%-26s' %('Node','JobID','Status','TimeIn','TimeUpdated')
    for task in tasks:
        print '%-10s\t%-16s\t%-10s\t%-26s\t%-26s' % (task.PartitionKey,task.RowKey,task.Status,task.TimeIn,task.TimeUpdated)
    return task

#read the queue and execute the jobs

starttime = datetime.now()

#loop continuously
while True:
    messages = queue_service.get_messages(qname,visibilitytimeout=maxtime)
    #if there is no work and more than allotted idle time has passed, shutdown
    if len(messages)==0:
        if (datetime.now() - starttime) > maxidle:
            #Note that you can replace this with a powershell command that shutdowns and deallocates VM
            call(["shutdown", "/s", "/t", "60"])
        sleep(10)
    else:
        #there is some work, so do it
        print len(messages)
        for message in messages:
            print (message.message_text)
            jobin = literal_eval(message.message_text)
            jobid=jobin['JobID']
            command=jobin['Command']
            filein=jobin['DataIn']
            fileout=jobin['DataOut']
            #check if the job was cancelled while waiting in queue
            task=jobquery(jobtable, jobid)
            if task.Status == 'Cancelled':
                print command, ' was cancelled before execution started.'
                queue_service.delete_message(qname, message.message_id, message.pop_receipt)
            #if not, carry on
            else:
                #update status table to running
                update_status(jobtable, node, jobid, task, 'Running')
                #Monitor commands for this job by subscribing to command topic and filtering
                #a node could run more than 1 job, hence have more than 1 subscription
                #16 is the max n. of cores today, one per job.
                subname='JobCommands-'+node+'-'+str(randint(1,16))
                print subname
                bus_service.create_subscription(cqname, subname)
                rule = Rule()
                rule.filter_type = 'SqlFilter'
                rule.filter_expression = 'JobID = '+'\''+jobid+'\''
                bus_service.delete_rule(cqname, subname, 'JobCommand', False)
                bus_service.create_rule(cqname, subname, 'JobCommand', rule, False)
                bus_service.delete_rule(cqname, subname, DEFAULT_RULE_NAME)

                #spawn process to execute command
                try:                    
                    p=Popen([command,filein,fileout])
                    returncode=p.poll()
                    #while the process is executing, listen to commands for it.
                    while (returncode == None):
                        try:
                            msg=bus_service.receive_subscription_message(cqname,subname)
                            print msg.body
                            if (msg.body == 'CANCEL'):
                                p.terminate()
                                msg.delete()
                                #On Unix subprocess.terminate() will return a negative code. On Windows, nothing.
                                returncode = -9
                            else:
                                sleep(5)
                                returncode=p.poll()
                            #other commands to be determined
                        except WindowsAzureError as error:
                            print 'No command yet...'
                            sleep(5)
                            returncode=p.poll()                                           

                    if returncode == 0:                  
                        #the job has run successfully - delete from queue and update status
                        queue_service.delete_message(qname, message.message_id, message.pop_receipt)
                        update_status(jobtable, node, jobid, task, 'Complete')
                        bus_service.delete_subscription(cqname,subname)
                        print command, ' executed successfully.'
                    elif returncode == -9:
                        #the job was cancelled - delete from queue and update status                     
                        queue_service.delete_message(qname, message.message_id, message.pop_receipt)
                        update_status(jobtable, node, jobid, task, 'Cancelled')
                        bus_service.delete_subscription(cqname,subname)
                        print command, ' was cancelled.'
                    else:
                        #the job has failed - delete from queue and update status                        
                        queue_service.delete_message(qname, message.message_id, message.pop_receipt)
                        update_status(jobtable, node, jobid, task, 'Failed')
                        bus_service.delete_subscription(cqname,subname)
                        print command, 'failed and returned ', returncode

                except OSError as error:
                    #the job has failed because of an os exception - remove from queue and update status                 
                    queue_service.delete_message(qname, message.message_id, message.pop_receipt)
                    update_status(jobtable, node, jobid, task, 'Failed')
                    bus_service.delete_subscription(cqname,subname)
                    print 'Execution failed: ', error

        starttime=datetime.now()
